#!/usr/bin/env python3
"""
Review Gate V3 - 独立MCP服务器（支持IPC通信）
作者: Lakshman Turlapati & daxian
通过Extension注册表发现Extension实例，使用IPC socket通信。

依赖要求:
- mcp>=1.9.2 (最新稳定版本)
- Python 3.8+
"""

import asyncio
import json
import sys
import logging
import os
import time
import uuid
import threading
import signal
import socket
import tempfile
import psutil
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional, Sequence, Set, List



from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    TextContent,
    Tool,
    ImageContent,
)

# 获取当前MCP进程的上下文信息
def get_mcp_context() -> Dict[str, Any]:
    """获取当前MCP进程的上下文信息"""
    try:
        current_process = psutil.Process()
        parent_process = current_process.parent()
        
        context = {
            'mcp_pid': current_process.pid,
            'parent_pid': parent_process.pid if parent_process else None,
        }
        
        logger.info(f"🔍 MCP进程上下文: PID={context['mcp_pid']}, 父PID={context['parent_pid']}")
        return context
    except Exception as e:
        logger.warning(f"⚠️ 获取MCP进程上下文失败: {e}")
        return {
            'mcp_pid': os.getpid(),
            'parent_pid': None,
        }

# 跨平台临时目录辅助函数
def get_temp_path(filename: str) -> str:
    """获取跨平台临时文件路径"""
    if os.name == 'nt':  # Windows
        temp_dir = tempfile.gettempdir()
    else:  # macOS和Linux
        temp_dir = '/tmp'
    return os.path.join(temp_dir, filename)

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]
)
# 设置日志路径
logger = logging.getLogger(__name__)
# 设置日志路径
log_file = get_temp_path('review_gate_v3_mcp.log')
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(file_handler)

class ExtensionDiscovery:
    """Extension发现和管理器"""
    
    def __init__(self):
        self.mcp_context = get_mcp_context()
        
    def discover_extensions_by_sockets(self) -> List[Dict]:
        """基于socket文件发现Extension实例（主要方法）"""
        try:
            import glob
            
            # 扫描/tmp目录下的review_gate_v3_*.sock文件
            socket_pattern = get_temp_path('review_gate_v3_*.sock')
            socket_files = glob.glob(socket_pattern)
            
            active_extensions = []
            
            for socket_path in socket_files:
                try:
                    # 从文件名提取PID
                    filename = os.path.basename(socket_path)
                    # 格式: review_gate_v3_{PID}.sock
                    if filename.startswith('review_gate_v3_') and filename.endswith('.sock'):
                        pid_str = filename[len('review_gate_v3_'):-len('.sock')]
                        extension_pid = int(pid_str)
                        
                        # 检查进程是否仍然存在
                        if self._is_process_alive(extension_pid):
                            extension_info = {
                                'extension_id': f'ext_socket_{extension_pid}_{int(time.time())}',
                                'extension_pid': extension_pid,
                                'ipc_path': socket_path,
                                'capabilities': ['chat', 'user_input'],
                                'version': '3.1.0',
                                'discovery_method': 'socket_scan',
                                'timestamp': os.path.getmtime(socket_path) * 1000  # 文件修改时间
                            }
                            active_extensions.append(extension_info)
                            logger.info(f"🔍 通过socket发现Extension: PID={extension_pid}, Socket={socket_path}")
                        else:
                            logger.warning(f"⚠️ Socket文件存在但进程已死亡: PID={extension_pid}, 清理socket文件")
                            try:
                                os.unlink(socket_path)
                            except:
                                pass
                                
                except (ValueError, OSError) as e:
                    logger.warning(f"⚠️ 解析socket文件失败: {socket_path}, 错误: {e}")
                    continue
            
            logger.info(f"🔍 通过socket扫描发现 {len(active_extensions)} 个活跃Extension")
            return active_extensions
            
        except Exception as e:
            logger.error(f"❌ 基于socket发现Extension失败: {e}")
            return []
    
    def _is_process_alive(self, pid: int) -> bool:
        """检查进程是否仍然存活"""
        try:
            import psutil
            return psutil.pid_exists(pid)
        except ImportError:
            # 如果没有psutil，使用系统调用
            try:
                import signal
                os.kill(pid, 0)
                return True
            except (OSError, ProcessLookupError):
                return False
        except:
            return False
        
    def discover_extensions(self) -> List[Dict]:
        """发现可用的Extension实例"""
        return self.discover_extensions_by_sockets()
    
    def find_best_extension(self) -> Optional[Dict]:
        """找到最佳的Extension实例 - 基于PID匹配"""
        extensions = self.discover_extensions()
        
        if not extensions:
            logger.warning("❌ 未找到任何活跃的Extension")
            return None
        
        # 基于PID匹配（同一窗口中extension_pid = parent_pid）
        parent_pid = self.mcp_context.get('parent_pid')
        if parent_pid:
            logger.info(f"🔍 基于PID匹配: {parent_pid}")
            for ext in extensions:
                ext_pid = ext.get('extension_pid')
                if ext_pid == parent_pid:
                    logger.info(f"✅ 找到PID匹配的Extension: {ext['extension_id']}")
                    return ext
        
        # 备用方案：选择最近活跃的Extension
        if extensions:
            latest_ext = max(extensions, key=lambda x: x.get('timestamp', 0))
            logger.warning(f"⚠️ 未找到PID匹配，使用最近活跃的Extension: {latest_ext['extension_id']}")
            return latest_ext
        
        return None

class ExtensionIPCClient:
    """Extension IPC客户端"""
    
    def __init__(self):
        self.socket = None
        self.connected = False
        
    async def connect_to_extension(self, extension_info: Dict) -> bool:
        """连接到Extension IPC服务器"""
        try:
            ipc_path = extension_info.get('ipc_path')
            if not ipc_path:
                logger.error("❌ Extension缺少IPC路径")
                return False
            
            # 创建socket连接
            self.socket = socket.socket(socket.AF_UNIX if os.name != 'nt' else socket.AF_INET, socket.SOCK_STREAM)
            
            if os.name == 'nt':
                # Windows命名管道处理
                # 这里需要特殊处理Windows命名管道
                logger.warning("⚠️ Windows命名管道连接需要特殊实现")
                return False
            else:
                # Unix域socket
                self.socket.connect(ipc_path)
            
            self.connected = True
            logger.info(f"✅ 已连接到Extension: {extension_info['extension_id']}")
            return True
            
        except Exception as e:
            logger.error(f"❌ 连接Extension失败: {e}")
            self.connected = False
            return False
    
    async def send_request(self, request_type: str, data: Dict, request_id: str = None) -> Optional[Dict]:
        """发送请求到Extension"""
        if not self.connected or not self.socket:
            logger.error("❌ 未连接到Extension")
            return None
        
        try:
            message = {
                'type': request_type,
                'id': request_id or str(uuid.uuid4()),
                'data': data,
                'timestamp': time.time()
            }
            
            # 发送消息
            message_json = json.dumps(message) + '\n'
            self.socket.send(message_json.encode('utf-8'))
            logger.info(f"📤 已发送请求到Extension: {request_type}")
            
            # 设置socket超时为0.5秒，减少等待时间
            self.socket.settimeout(0.5)
            
            # 等待响应
            response_data = b''
            timeout_counter = 0
            max_timeout_seconds = 900  # 15分钟
            max_timeout_counter = max_timeout_seconds * 2  # 每0.5秒一次，所以需要 * 2
            
            while timeout_counter < max_timeout_counter:
                try:
                    chunk = self.socket.recv(1024)
                    if not chunk:
                        logger.warning("⚠️ Extension连接已关闭")
                        break
                    response_data += chunk
                    if b'\n' in response_data:
                        logger.info("📨 收到完整响应数据")
                        break
                except socket.timeout:
                    timeout_counter += 1  # socket超时0.5秒后增加1，所以每0.5秒增加1
                    # 移除额外的sleep，避免双重延时
                    if timeout_counter % 60 == 0:  # 每30秒记录一次（60 * 0.5s = 30s）
                        actual_seconds = timeout_counter * 0.5
                        logger.info(f"⏳ 等待Extension响应中... ({actual_seconds:.1f}s / {max_timeout_seconds}s)")
                    continue
                except Exception as e:
                    logger.error(f"❌ 接收数据时出错: {e}")
                    break
            
            # 检查是否超时（转换为实际秒数）
            actual_timeout_seconds = timeout_counter * 0.5
            if actual_timeout_seconds >= max_timeout_seconds:
                logger.warning(f"⚠️ 请求超时 ({actual_timeout_seconds:.1f}秒)")
                return {
                    'type': 'timeout',
                    'error': f'请求在{actual_timeout_seconds:.1f}秒后超时',
                    'timeout_duration': actual_timeout_seconds
                }
            
            if response_data:
                try:
                    response_text = response_data.decode('utf-8').strip()
                    logger.info(f"📨 原始响应数据: {response_text[:200]}...")
                    response = json.loads(response_text)
                    logger.info(f"📨 收到Extension响应: {response.get('type', 'unknown')}")
                    return response
                except json.JSONDecodeError as e:
                    logger.error(f"❌ 响应JSON解析失败: {e}")
                    logger.error(f"原始数据: {response_data[:500]}")
            else:
                logger.warning("⚠️ 未收到响应数据")
            
        except Exception as e:
            logger.error(f"❌ 发送请求失败: {e}")
            self.connected = False
        
        return None
    
    def disconnect(self):
        """断开连接"""
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
            self.socket = None
        self.connected = False

class ReviewGateV3Server:
    """Review Gate V3 独立MCP服务器类"""
    
    def __init__(self):
        logger.info("🚀 正在初始化 Review Gate V3 独立MCP服务器...")
        
        self.server = Server("review-gate-v3-standalone")
        self.setup_handlers()
        self.shutdown_requested = False
        self.shutdown_reason = ""
        
        # Extension发现和IPC客户端
        self.extension_discovery = ExtensionDiscovery()
        self.ipc_client = ExtensionIPCClient()
        
        logger.info("✅ Review Gate V3 独立MCP服务器基础初始化完成")

    def setup_handlers(self):
        """设置MCP请求处理器"""
        
        @self.server.list_tools()
        async def list_tools():
            """列出Cursor Agent可用的Review Gate工具"""
            logger.info("🔧 Cursor Agent请求可用工具列表")
            tools = [
                Tool(
                    name="review_gate_chat",
                    description="在Cursor中打开Review Gate聊天弹窗进行反馈和审核。通过Extension注册表发现Extension实例，使用IPC socket通信。当您需要用户输入、反馈或人类用户审核时使用此工具。",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "message": {
                                "type": "string",
                                "description": "在Review Gate弹窗中显示的消息 - 这是用户将看到的内容",
                                "default": "请提供您的审核或反馈："
                            },
                            "title": {
                                "type": "string", 
                                "description": "Review Gate弹窗窗口的标题",
                                "default": "Review Gate V3"
                            },
                            "context": {
                                "type": "string",
                                "description": "关于需要审核内容的额外上下文信息（代码、实现等）",
                                "default": ""
                            },
                            "timeout": {
                                "type": "integer",
                                "description": "等待用户响应的超时时间（秒）",
                                "default": 900
                            }
                        }
                    }
                )
            ]
            logger.info(f"✅ 为Cursor Agent列出了{len(tools)}个Review Gate工具")
            return tools

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict):
            """处理来自Cursor Agent的工具调用"""
            logger.info(f"🎯 CURSOR AGENT调用工具: {name}")
            logger.info(f"📋 工具参数: {arguments}")
            
            try:
                if name == "review_gate_chat":
                    return await self._handle_review_gate_chat(arguments)
                else:
                    logger.error(f"❌ 未知工具: {name}")
                    raise ValueError(f"未知工具: {name}")
                    
            except Exception as e:
                logger.error(f"💥 工具调用错误 {name}: {e}")
                return [TextContent(type="text", text=f"错误: 工具 {name} 执行失败: {str(e)}")]

    async def _handle_review_gate_chat(self, args: dict) -> list:
        """处理Review Gate聊天弹窗并等待用户输入"""
        message = args.get("message", "请提供您的审核或反馈：")
        title = args.get("title", "Review Gate V3")
        context = args.get("context", "")
        timeout = args.get("timeout", 900)
        
        logger.info(f"🎯 开始处理review_gate_chat请求")
        logger.info(f"📝 请求参数: message='{message[:50]}...', timeout={timeout}s")
        
        # 发现并连接到最佳Extension
        best_extension = self.extension_discovery.find_best_extension()
        if not best_extension:
            error_msg = "错误: 未找到可用的Extension实例。请确保Cursor Extension正在运行。"
            logger.error(f"❌ {error_msg}")
            
            # 显示调试信息
            all_extensions = self.extension_discovery.discover_extensions()
            if all_extensions:
                logger.error("📋 当前注册的Extension列表:")
                for i, ext in enumerate(all_extensions):
                    logger.error(f"   {i+1}. {ext.get('extension_id', 'unknown')}")
                    logger.error(f"      Extension PID: {ext.get('extension_pid', 'unknown')}")
                    logger.error(f"      IPC路径: {ext.get('ipc_path', 'unknown')}")
            
            return [TextContent(type="text", text=error_msg)]
        
        logger.info(f"✅ 选择的Extension详情:")
        logger.info(f"   Extension ID: {best_extension.get('extension_id', 'unknown')}")
        logger.info(f"   Extension PID: {best_extension.get('extension_pid', 'unknown')}")
        logger.info(f"   IPC路径: {best_extension.get('ipc_path', 'unknown')}")
        
        # 连接到Extension
        connected = await self.ipc_client.connect_to_extension(best_extension)
        if not connected:
            error_msg = f"错误: 无法连接到Extension {best_extension['extension_id']}"
            logger.error(f"❌ {error_msg}")
            return [TextContent(type="text", text=error_msg)]
        
        try:
            # 发送用户输入请求
            request_id = str(uuid.uuid4())
            request_data = {
                'message_id': request_id,
                'message': message,
                            'title': title,
            'context': context,
            'timeout': timeout
            }
            
            logger.info(f"📤 向Extension发送用户输入请求: {request_id}")
            
            # 发送请求并等待响应
            response = await self.ipc_client.send_request('request_user_input', request_data, request_id)
            
            # 处理超时响应
            if response and response.get('type') == 'timeout':
                timeout_duration = response.get('timeout_duration', timeout)
                error_msg = f"⏰ 超时: 在{timeout_duration}秒内未收到用户输入。请检查Review Gate弹窗是否正常显示。"
                logger.warning(f"⚠️ {error_msg}")
                return [TextContent(type="text", text=error_msg)]
            
            if response and response.get('type') == 'response':
                response_data = response.get('data', {})
                user_input = response_data.get('user_input', '')
                attachments = response_data.get('attachments', [])
                
                logger.info(f"✅ 收到用户响应: {user_input[:50]}...")
                
                response_content = [TextContent(type="text", text=f"用户响应: {user_input}")]
                
                # 处理附件
                if attachments:
                    logger.info(f"📎 收到 {len(attachments)} 个附件")
                    for i, attachment in enumerate(attachments):
                        logger.info(f"📎 附件 {i+1} 字段: {list(attachment.keys())}")
                        
                        # 支持多种字段名格式以兼容不同版本
                        mime_type = (attachment.get('mimeType') or 
                                   attachment.get('mime_type') or 
                                   attachment.get('type', ''))
                        
                        base64_data = (attachment.get('base64Data') or 
                                     attachment.get('data') or 
                                     attachment.get('base64'))
                        
                        file_name = (attachment.get('fileName') or 
                                   attachment.get('name') or 
                                   f'image_{i+1}')
                        
                        logger.info(f"📎 解析结果: mime_type={mime_type}, has_data={bool(base64_data)}, file_name={file_name}")
                        
                        if mime_type.startswith('image/') and base64_data:
                            try:
                                # 处理data URL格式的Base64数据
                                if base64_data.startswith('data:'):
                                    # 提取纯Base64数据
                                    parts = base64_data.split(',')
                                    if len(parts) == 2:
                                        base64_data = parts[1]
                                        # 从data URL header中提取MIME类型
                                        header = parts[0]
                                        if 'image/' in header:
                                            mime_type = header.split(';')[0].replace('data:', '')
                                
                                image_content = ImageContent(
                                    type="image",
                                    data=base64_data,
                                    mimeType=mime_type
                                )
                                response_content.append(image_content)
                                logger.info(f"✅ 成功添加图片到响应: {file_name} ({mime_type})")
                            except Exception as e:
                                logger.error(f"❌ 添加图片到响应时出错: {e}")
                                logger.error(f"❌ 附件详情: {attachment}")
                        else:
                            logger.warning(f"⚠️ 跳过非图片附件或无效数据: mime_type={mime_type}, has_data={bool(base64_data)}")
                else:
                    logger.info("📎 没有收到附件")
                
                return response_content
            else:
                error_msg = f"超时: 在{timeout}秒内未收到用户输入"
                logger.warning(f"⚠️ {error_msg}")
                return [TextContent(type="text", text=error_msg)]
                
        finally:
            # 断开连接
            self.ipc_client.disconnect()

    async def run(self):
        """运行Review Gate V3 独立MCP服务器"""
        logger.info("🚀 启动Review Gate V3 独立MCP服务器...")
        
        try:
            # 启动MCP服务器
            async with stdio_server() as (read_stream, write_stream):
                logger.info("✅ Review Gate V3 独立服务器在stdio传输上活跃")
                
                await self.server.run(
                    read_stream,
                    write_stream,
                    self.server.create_initialization_options()
                )
        except Exception as e:
            logger.error(f"❌ 服务器运行时错误: {e}")
            raise
        finally:
            self.shutdown_requested = True

def setup_signal_handlers(server):
    """设置信号处理器"""
    def signal_handler(signum, frame):
        logger.info(f"🛑 收到信号 {signum}，正在优雅关闭...")
        server.shutdown_requested = True
        sys.exit(0)
    
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

async def main():
    """Review Gate V3 独立MCP服务器的主入口点"""
    logger.info("🎬 启动Review Gate V3 独立MCP服务器...")
    logger.info(f"Python版本: {sys.version}")
    logger.info(f"平台: {sys.platform}")
    logger.info(f"操作系统: {os.name}")
    
    try:
        server = ReviewGateV3Server()
        setup_signal_handlers(server)
        await server.run()
    except Exception as e:
        logger.error(f"❌ MCP服务器致命错误: {e}")
        import traceback
        logger.error(traceback.format_exc())
        raise

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("🛑 用户停止服务器")
    except Exception as e:
        logger.error(f"❌ 服务器崩溃: {e}")
        sys.exit(1) 